from network_topology import weighted_dijkstra
import pandas as pd
import networkx as nx
import numpy as np
import statistics
from const import TASK_NUM

'''
尝试使用evaluate评估一个环境下的奖励，这个东西的和就是累计奖励和
图的拓扑结构和权重包含在图对象中，任务和起始终止包含在task_dic中
（graph, weight, task）构成一个完整的环境Si
G:graph 
task_dic:{'tasks':..., 'pairs':...}
capacity:pd.dataframe
'''


def softmax(x):
    return np.exp(x) / np.sum(np.exp(x))


# 这里传入的是task_generator的第二个返回值，输出的的是一个环境对应的奖励ri，对应一组任务
def evaluate(G, task_dic, capacity_matrix, v_matrix):
    paths = weighted_dijkstra(G)

    '''
    Warning!!!
    1、要解决奖励过大的问题，需要对奖励的每个部分进行分析
    2、尝试：将预测权值的好坏也加入evaluate之中，权值的变化不能过大过小
    
    下面的是应当确定根据哪些情况计算奖励
    1、每条路线不能超载
    2、传输延时不能太大
    3、...
    '''

    delay_list = []
    payload_matrix = np.zeros_like(capacity_matrix)
    capacity_reward = 0
    delay_reward = 0

    for key, value in task_dic.items():
        source = key[0]
        target = key[1]
        path = paths[source][1][target]  # type = list
        # 这里path是排过序的，所以source序号和在list中的序号一致
        # 因为目标点被作为key，所以转成字符串
        # 每个经过的路径上都要加上这个payload

        total_delay = 0

        for i in range(0, path.__len__() - 1):

            a = path[i]
            b = path[i + 1]
            payload_matrix[a][b] += value
            payload_matrix[b][a] += value
            total_delay += np.log(value / v_matrix[a][b]) / np.log(v_matrix[a][b])

        delay_list.append(total_delay)

    # 对于每条路径：
    # 负载<=容量 reward+=1
    # 负载>容量  reward-=1
    compare_matrix = np.less_equal(payload_matrix, capacity_matrix)
    for row in compare_matrix:
        for le in row:
            if le:
                capacity_reward += 1
            else:
                capacity_reward -= 5
    # capacity_reward = np.log(capacity_reward)

    # 每个delay_list应该是有十个元素，正好存了10个任务计算出来的传输时延，应该越小给的奖励越多，但是这个量级应该和capacity_reward匹配
    # 先尝试取对数，时延越小奖励增加幅度越大

    for item in delay_list:
        delay_reward -= item
    delay_reward = max(-40, delay_reward)
    delay_reward = min(40, delay_reward)
    # delay_reward = np.log(delay_reward)

    return capacity_reward / 10 + delay_reward / 10


# 针对的应该是group,一个group内部对应着若干次神经网络对权值的预测
def accumulated_reward(G, task_group, capacity_matrix, v_matrix, alpha):
    group_num = task_group.__len__()
    reward_list = []
    for i in range(0, group_num):
        task_dic = task_group[i]['task_dic']
        reward = evaluate(G, task_dic, capacity_matrix, v_matrix)
        reward_list.append(reward)
    index = group_num - 2
    while index > 0:
        reward_list[index] = reward_list[index] + alpha * reward_list[index + 1]  # 求出累计奖励和，奖励预测神经网络应该管这块
    return reward_list


# 针对的应该是group,一个group内部对应着若干次神经网络对权值的预测
def group_reward(G, task_group, capacity_matrix, v_matrix):
    group_num = task_group.__len__()
    reward_list = []
    for i in range(0, group_num):
        task_dic = task_group[i]['task_dic']
        reward = evaluate(G, task_dic, capacity_matrix, v_matrix)
        reward_list.append(reward)

    return reward_list
